#pragma once

#include <NvInfer.h>
#include <iostream>
#include <memory>
#include <opencv2/core.hpp>
#include <rclcpp/rclcpp.hpp>

namespace perception::camera {

class Logger : public nvinfer1::ILogger {
  public:
    explicit Logger(Severity severity = Severity::kWARNING)
        : reportable_severity(severity) {}

    void log(Severity severity, const char* msg) noexcept {
        if (severity > reportable_severity) {
            return;
        }
        switch (severity) {
            case Severity::kINTERNAL_ERROR:
                std::cerr << "lane INTERNAL_ERROR: ";
                break;
            case Severity::kERROR:
                std::cerr << "lane ERROR: ";
                break;
            case Severity::kWARNING:
                std::cerr << "lane WARNING: ";
                break;
            case Severity::kINFO:
                std::cerr << "lane INFO: ";
                break;
            default:
                std::cerr << "lane UNKNOWN: ";
                break;
        }
        std::cerr << msg << std::endl;
    }
    Severity reportable_severity;
};

struct Detection {
    float unknown;
    float score;
    float start_y;
    float start_x;
    float length;
    float lane_xs[72];
};

class LaneDetector {
  public:
    bool init(const std::string& model_dir);
    void deinit();
    void process(const cv::Mat& raw_image, std::vector<std::vector<cv::Point2f>> &lanes);

    static std::shared_ptr<LaneDetector> get_detector(const std::string& model_dir);


  private:
    LaneDetector();
    // ~LaneDetector();
    void LoadEngine(const std::string& engine_file);

    void PostProcess(std::vector<std::vector<cv::Point2f>> &lanes, float conf_thresh=0.4f, float nms_thresh=50.f, int nms_topk=4);
    float LaneIoU(const Detection& a, const Detection& b);
    Logger g_logger_;
    rclcpp::Logger logger__;

    cudaStream_t stream_;
    nvinfer1::ICudaEngine* engine_;
    nvinfer1::IExecutionContext* context_;
    
    void* buffers_[2];
    int buffer_size_[2];
    std::vector<float> image_data_;
    std::vector<Detection> detections_;
    static std::shared_ptr<LaneDetector> detector_;
    int orig_img_h_ = 1080;
    int orig_img_w_ = 1920;
    const int INPUT_H_ = 360;
    const int INPUT_W_ = 640;
    const int N_OFFSETS_ = 72;
    const int N_STRIPS_ = (72 - 1);
    const int MAX_COL_BLOCKS_ = 1000;
};
}